package httpstunnel

import (
	"bufio"
	"context"
	"encoding/binary"
	"errors"
	"github.com/go-acme/lego/log"
	"io"
	"net"
	"net/http"
	"sync"
)

var (
	tcpListener net.Listener
)

func runTcp(addr string) error {
	ln, err := net.Listen("tcp", addr)
	if err != nil {
		return err
	}
	tcpListener = ln
	go acceptTcp()
	return nil
}

func acceptTcp() {
	for {
		conn, err := tcpListener.Accept()
		if err != nil {
			log.Println("[ERROR] ", err)
			continue
		}

		go handleConn(conn)
	}
}

func handleConn(conn net.Conn) {
	defer func() {
		conn.Close()
		registryDomain.Range(func(key, value interface{}) bool {
			if value.(net.Conn) == conn {
				registryDomain.Delete(key)
			}
			return true
		})
	}()
	for {
		var b = make([]byte, 2)
		_, err := io.ReadFull(conn, b)
		if err != nil {
			log.Println("[ ERROR ] connection read error ", err)
			return
		}
		cmdLen := binary.BigEndian.Uint16(b)

		var cmd = make([]byte, cmdLen)
		_, err = io.ReadFull(conn, cmd)
		if err != nil {
			log.Println("[ ERROR ] connection read error ", err)
			return
		}

		bodyHeadLen := make([]byte, 2)
		_, err = io.ReadFull(conn, bodyHeadLen)
		if err != nil {
			log.Println("[ ERROR ] connection read error ", err)
			return
		}

		bh := binary.BigEndian.Uint16(bodyHeadLen)
		var body = make([]byte, bh)

		_, err = io.ReadFull(conn, body)
		if err != nil {
			log.Println("[ ERROR ] connection read error ", err)
			return
		}
		switch string(cmd) {
		case CmdRegistry:
			err = handleRegistry(conn, body)
			if err != nil {
				log.Println("[ ERROR ] registry error ", err)
				return
			}
		case CmdDeRegistry:
			err = handleDeRegistry(conn, body)
		case CmdProxy:
			idx := body[len(body)-4:]
			w, ok := respWriter.Load(idx)
			if ok {
				respWriter.Delete(idx)
				ctx := w.(context.Context)
				wr := ctx.Value("w").(http.ResponseWriter)
				ch := ctx.Value("chan").(chan struct{})
				wr.Write(body[len(body)-4:])
				ch <- struct{}{}
			}
		case CmdPing:
			_, _ = conn.Write(tcpWriteWrap(CmdPong, []byte("pong")))
		case CmdPong:
			_, _ = conn.Write(tcpWriteWrap(CmdPing, []byte("ping")))
		}
	}
}

// -----------------------------
//   2   |   cmd   |  2  |  body
// -----------------------------
const (
	CmdRegistry   = "register"
	CmdDeRegistry = "deregister"
	CmdProxy      = "proxy"
	CmdPing       = "ping"
	CmdPong       = "pong"
	CmdError      = "error"
)

func tcpWriteWrap(cmd string, body []byte) []byte {
	var l = 2 + len(cmd) + 2 + len(body)
	b := make([]byte, l)
	h := len(cmd)
	binary.BigEndian.PutUint16(b[:2], uint16(h))
	copy(b[2:2+len(cmd)], []byte(cmd))
	bl := len(body)
	binary.BigEndian.PutUint16(b[2+len(cmd):2+len(cmd)+2], uint16(bl))
	copy(b[2+len(cmd)+2:], body)

	bufio.NewReader(clientConn).ReadLine()

	return b
}

var (
	registryDomain = sync.Map{}
)

func handleRegistry(conn net.Conn, body []byte) error {
	d := string(body)
	_, ok := registryDomain.Load(d)
	if ok {
		conn.Write(tcpWriteWrap(CmdError, []byte("domain already exists")))
		return errors.New("domain already exists")
	}
	registryDomain.Store(d, conn)
	return nil
}

func handleDeRegistry(conn net.Conn, body []byte) error {
	registryDomain.Delete(string(body))
	return nil
}
